import pandas as pd

from scipy.sparse import csr_matrix
import cupy as cp

def build_matrix(data):
    # 确保 user 和 item 是整数类型
    # 使用 pandas.factorize 对 user 和 item 进行整数编码
    data['user'], user_uniques = pd.factorize(data['user'])
    data['item'], item_uniques = pd.factorize(data['item'])
    
    # 将评分转换为整数
    data['rating'] = data['rating'].astype(int)
    ratings_matrix = csr_matrix((data['rating'], (data['user'], data['item'])))
    return ratings_matrix

# 使用CuPy代替NumPy进行矩阵操作
def als_wr_gpu(ratings_matrix, num_features=10, lambda_reg=0.1, num_iterations=50, early_stopping_rounds=5):
    print("starting als_wr...")
    num_users, num_items = ratings_matrix.shape
    user_features = cp.random.rand(num_users, num_features)
    item_features = cp.random.rand(num_items, num_features)
    
    best_loss = float('inf')
    epochs_without_improvement = 0  # 记录连续未改进的轮次

    for i in range(num_iterations):
        print(f"als迭代 {i}")
        
        # 更新用户特征矩阵
        for u in range(num_users):
            print(f"u:::{u}")
            idx = ratings_matrix[u].indices
            R_u = cp.asarray(ratings_matrix[u, idx].toarray().flatten())  # 转换为cupy.ndarray
            X_u = item_features[idx, :]
            user_features[u, :] = cp.linalg.solve(cp.dot(X_u.T, X_u) + lambda_reg * cp.eye(num_features), cp.dot(X_u.T, R_u))
        
        # 更新物品特征矩阵
        for j in range(num_items):
            print(f"i:::{j}")
            idx = ratings_matrix[:, j].indices
            R_i = cp.asarray(ratings_matrix[idx, j].toarray().flatten())  # 转换为cupy.ndarray
            X_i = user_features[idx, :]
            item_features[j, :] = cp.linalg.solve(cp.dot(X_i.T, X_i) + lambda_reg * cp.eye(num_features), cp.dot(X_i.T, R_i))
        
        # 计算当前损失
        loss = 0
        for u in range(num_users):
            idx = ratings_matrix[u].indices
            R_u = cp.asarray(ratings_matrix[u, idx].toarray().flatten())  # 转换为cupy.ndarray
            X_u = item_features[idx, :]
            loss += cp.sum((R_u - cp.dot(X_u, user_features[u, :].T)) ** 2)
        
        for j in range(num_items):
            idx = ratings_matrix[:, j].indices
            R_i = cp.asarray(ratings_matrix[idx, j].toarray().flatten())  # 转换为cupy.ndarray
            X_i = user_features[idx, :]
            loss += cp.sum((R_i - cp.dot(X_i, item_features[j, :].T)) ** 2)
        
        # 加上正则化项
        loss += lambda_reg * (cp.sum(user_features ** 2) + cp.sum(item_features ** 2))
        
        # 检查是否需要提前停止
        if loss < best_loss:
            best_loss = loss
            epochs_without_improvement = 0
        else:
            epochs_without_improvement += 1
        
        if epochs_without_improvement >= early_stopping_rounds:
            print(f"提前停止：在 {i} 轮迭代后没有显著改善。")
            break

    return user_features, item_features